import torch
import torch.nn as nn


class ConvTemporalGraphical(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 t_kernel_size=1,
                 t_stride=1,
                 t_padding=0,
                 t_dilation=1,
                 bias=True):
        super().__init__()
        # 卷积核的尺寸（这里的kernel_size一般是空间维度的核大小）
        self.kernel_size = kernel_size
        # 定义一个2D卷积层，用于处理输入数据
        # in_channels: 输入通道数
        # out_channels * kernel_size: 输出通道数（这里乘以kernel_size是为了后续处理）
        # kernel_size=(t_kernel_size, 1): 卷积核的时间和空间尺寸
        # padding=(t_padding, 0): 时间和空间维度的填充
        # stride=(t_stride, 1): 时间和空间维度的步长
        # dilation=(t_dilation, 1): 时间和空间维度的膨胀
        # bias=bias: 是否使用偏置
        self.conv = nn.Conv2d(
            in_channels,
            out_channels * kernel_size,
            kernel_size=(t_kernel_size, 1),
            padding=(t_padding, 0),
            stride=(t_stride, 1),
            dilation=(t_dilation, 1),
            bias=bias)

    def forward(self, x, A):
        # 断言输入的A的尺寸的第0维等于卷积核尺寸kernel_size
        assert A.size(0) == self.kernel_size

        # 对输入x应用卷积操作
        x = self.conv(x)

        # 获取卷积输出的尺寸信息，n: 批量大小, kc: 卷积后的通道数, t: 时间维度大小, v: 空间维度大小
        n, kc, t, v = x.size()
        # 调整x的形状，方便后续与邻接矩阵A进行运算
        x = x.view(n, self.kernel_size, kc // self.kernel_size, t, v)
        # 使用爱因斯坦求和约定（einsum）对x和A进行运算，这里的操作是为了融合空间信息
        x = torch.einsum('nkctv,kvw->nctw', (x, A))

        # 返回连续的张量，确保内存中数据的连续性，以提高计算效率
        return x.contiguous()
